import baostock as bs
import pandas as pd
import mplfinance as mpf

class DefTypesPool():
    def __init__(self):
        self.routes = {}
    # 注册绘图函数
    def route_types(self, types_str):
        def decorator(f):
            self.routes[types_str] = f
            return f
        return decorator
    # 发现绘图函数
    def route_output(self, path):
        function_val = self.routes.get(path)
        if function_val:
            return function_val
        else:
            raise ValueError('Route "{}"" has not been registered'.format(path))
        
class MplTypesDraw():
    mpl = DefTypesPool() 
                          

class MplVisualIf(MplTypesDraw):
    def __init__(self):
        MplTypesDraw.__init__(self)
    


    
    def fig_output2(self, **kwargs):
        self.index =kwargs['index']
        df_dat = kwargs['draw_kind']['ochl']
        s = mpf.make_mpf_style(base_mpf_style='yahoo', rc={'font.family': 'SimHei'})
        pdata = pd.DataFrame.from_dict(df_dat)
 
        apds = []
        if kwargs['draw_kind'].__contains__('annotate2'):
            annotate2 = kwargs['draw_kind']['annotate2']
            maxdf = annotate2['max']['andata']
            mindf = annotate2['min']['andata']
            print(maxdf)
            apds = [mpf.make_addplot(pdata,type='scatter',markersize=200,marker='^'),mpf.make_addplot(pdata,type='scatter',markersize=200,marker='v')]
        mpf.plot( data = pdata, addplot=apds, type='candle',title = u'k线图' ,style = s, ylabel = u'价格', xlabel = u'日期',figratio=(23,8),figscale=0.85)
        


# 登陆系统
def bs_k_data_stock(code_val='sz.000651',start_val = '2025-01-01', end_val='2025-05-01',freq_val='d',adjust_val='3'):
    lg = bs.login()
    # 获取历史行情数据
    fields = "date,open,high,low,close,volume"
    df_bs = bs.query_history_k_data_plus(code_val, fields, start_date=start_val, end_date=end_val,frequency=freq_val ,adjustflag=adjust_val)
    #日K线，3，默认不赋权，1 后复权，2 前复权
    data_list= []
    while (df_bs.error_code == '0') & df_bs.next():
        #获取一条记录，将记录合并在一起
        data_list.append(df_bs.get_row_data())
    result = pd.DataFrame(data_list, columns=df_bs.fields)

    result.close = result.close.astype('float64')
    result.open = result.open.astype('float64')
    result.low = result.low.astype('float64')
    result.high = result.high.astype('float64')
    result.volume = result.volume.astype('float64')
    result.volume = result.volume/100 #单位转换:股-手
    result.date = pd.DatetimeIndex(result.date)
    result.set_index("date", drop=True, inplace= True)
    result.index = result.index.set_names('Date')

    recon_data = {'High' : result.high, 'Low': result.low, 'Open': result.open, 'Close':result.close, 'Volume': result.volume}
    df_recon = pd.DataFrame(recon_data)

    #退出系统
    bs.logout()
    return df_recon

app = MplVisualIf()

# 周线
def draw_kweek_chart(stock_dat):
    # 周期重采样
    # rule='W'周how=last()最后一天closed='right'右闭合label='right'右标签
    Freq_T = 'W-FRI'

    # 周线Close等于一周中最后一个交易日Close
    week_dat = stock_dat.resample(Freq_T, closed='right',label='right').last()
    # 周线Open等于一周中第一个交易日Close
    week_dat.Open = stock_dat.Open.resample(Freq_T, closed='right', label='right').first()
    # 周线High等于一周中High的最大值
    week_dat.High = stock_dat.High.resample(Freq_T, closed='right', label='right').max()
    # 周线Low等于一周中Low的最小值
    week_dat.Low = stock_dat.Volume.resample(Freq_T, closed='right',label='right').min()
    # 周线Volume等于一周中Volume的总和
    week_dat.Volume = stock_dat.Volume.resample(Freq_T, closed='right',label='right').sum()
    print(week_dat.head())

    layout_dict = {'figsize':(14,7),
                   'index': week_dat.index,
                   'draw_kind': {
                       'ochl':
                       {
                           'Open':week_dat.Open,
                           'Close':week_dat.Close,
                           'High':week_dat.High,
                           'Low':week_dat.Low
                       }
                   },
                   'title': u'601595-周k线',
                   'ylabel':u"价格"
                   }
    app.fig_output2(**layout_dict)

def draw_fibonacci_chart(stock_dat):
    # 黄金分割线
    Fib_max = stock_dat.Close.max()
    Fib_maxid = stock_dat.index.get_loc(stock_dat.Close.idxmax())
    Fib_min = stock_dat.Close.min()
    Fib_minid = stock_dat.index.get_loc(stock_dat.Close.idxmin())
    Fib_382 = (Fib_max - Fib_min) * 0.382 + Fib_min
    Fib_618 = (Fib_max - Fib_min) * 0.618 + Fib_min
    print(u'黄金分割0.382:{}'.format(round(Fib_382, 2)))
    print(u'黄金分割0.618:{}'.format(round(Fib_618, 2)))
    max_df = stock_dat[stock_dat.Close == stock_dat.Close.max()]
    min_df = stock_dat[stock_dat.Close == stock_dat.Close.min()]

    #绘制K线图+支撑/阻力
    layout_dict = {'figsize':(14,7),
                   'index':stock_dat.index,
                   'draw_kind':{
                       'ochl':{
                           'Open':stock_dat.Open,
                           'Close':stock_dat.Close,
                           'High':stock_dat.High,
                           'Low':stock_dat.Low
                       },
                       'hline2':
                       {
                           'Fib_382':{
                               'pos':Fib_382,
                               'c': 'r'
                           },
                           'Fib_618':{
                               'pos':Fib_618,
                               'c':'g'
                           }
                       },
                       'annotate2':{
                           u'max':{
                               'andata':max_df ,
                               'va':'top',
                               'xy_y':'High',
                               'xytext':(-30, stock_dat.Close.mean()),
                               'fontsize': 8,
                               'arrow':dict(facecolor='red', shrink = 0.1)
                           },
                           u'min':{
                               'andata':min_df,
                               'va':'bottom',
                               'xy_y':'Low',
                               'xytext':(-30, -stock_dat.Close.mean()),
                               'fontsize': 8,
                               'arrow':dict(facecolor='green', shrink = 0.1)
                           }
                       }
                   },
                   'title':'601594 支撑/阻力位',
                   'ylabel':u'价格',
                   'legend':u'best'}
    app.fig_output2(**layout_dict)


df_stockload = bs_k_data_stock('sh.601595','2024-07-28', '2025-07-04')
draw_fibonacci_chart(df_stockload)
draw_kweek_chart(df_stockload)
